

import numpy as np
import os
import h5py

fold_list_name = '/dataset/Kaggle/fold_list.npy'
# label_name = '/data-disk/Kaggle/lidc&kaggle/label_dict3.npy'
# proj_dir = ''
CANDIS_DIR = 'output/candis_60000'
# CANDIS_DIR = ''
CANDIS_LUNGMASK_DIR = 'output/candis_before_nms_lungmask_60000'

thresh_val = 0.98

def filter_lungmask():
    fold_list = np.load(fold_list_name)
    labels_kaggle = np.load('/data-disk/Kaggle/backup/kaggle/label_dict_0325.npy').item()
    labels_lidc = np.load('/data-disk/Kaggle/backup/lidc/label_dict_3_0329.npy').item()
    labels = dict(labels_kaggle.items() + labels_lidc.items())
    lungmask_kaggle = h5py.File('/data-disk/Kaggle/backup/kaggle/lungmask.hdf5', 'r')
    lungmask_lidc = h5py.File('/data-disk/Kaggle/backup/lidc/lungmask.hdf5', 'r')

    dataset_kaggle = h5py.File('/data-disk/Kaggle/backup/kaggle/vol.hdf5', 'r')
    dataset_lidc = h5py.File('/data-disk/Kaggle/backup/lidc/vol.hdf5', 'r')

    is_found_all = []
    is_mali_all = []
    dict_not_found = {}
    cnt_fp = 0
    cnt_tp = 0
    iscan = 0

    for fold_idx in range(1, 6):
        scan_names = fold_list[fold_idx - 1]
        # proj_dir = '/ssd1/Kaggle/python/faster_rcnn/anchor_fold' + str(fold_idx)
        proj_dir = '/ssd1/dingj/kaggle/faster_rcnn/all_nodule_flip_ad/anchor_fold' + str(
            fold_idx) + '_all_flip_ad'
        if not os.path.exists(os.path.join(proj_dir, CANDIS_LUNGMASK_DIR)):
            os.mkdir(os.path.join(proj_dir, CANDIS_LUNGMASK_DIR))

        for scan_name in scan_names:
            iscan = iscan + 1
            # if scan_name[0] == 'L':
            #     continue
            if scan_name[0] == 'L':
                spacing = dataset_lidc[scan_name].attrs['spacing'][0:2]
            else:
                spacing = dataset_kaggle[scan_name].attrs['spacing'][0:2]
            candis = np.load(os.path.join(proj_dir, CANDIS_DIR, scan_name, 'candis.npy'))
            label = labels[scan_name]
            is_found = np.zeros((len(label), 1))
            is_mali = np.zeros((len(label), 1))

            if scan_name[0] == 'L':
                lungmask = lungmask_lidc[scan_name][:]
            else:
                lungmask = lungmask_kaggle[scan_name][:]

            for ilabel in xrange(len(label)):
                nodule = label[ilabel]
                if scan_name[0] == 'L':
                    if nodule['attrs']['malignancy'] > 3:
                        is_mali[ilabel] = 1
                if scan_name[0] != 'L':
                    if nodule['malign'] > 6:
                        is_mali[ilabel] = 1
            is_mali_all.extend(is_mali)
            for icandi in xrange(len(candis)):  # for every slice
                candi_slice = candis[icandi]
                if len(candi_slice) == 0:
                    continue

                _candi_slice = []
                for bbox in candi_slice:  # for every candidate in one slice
                    # if bbox[4] < thresh_val:
                    #     continue
                    if np.max([(bbox[3] - bbox[1]),
                               (bbox[2] - bbox[0])] * spacing) < 5:  # drop bbox with size less than 5mm
                        continue
                    candi_r = max([(bbox[3] - bbox[1]) / 2, (bbox[2] - bbox[0]) / 2])
                    candi_center = [(bbox[1] + bbox[3]) / 2, (bbox[0] + bbox[2]) / 2, icandi]
                    if candi_r < 10:
                        if not lungmask[candi_center[0], candi_center[1], candi_center[2]]:
                            continue
                    _candi_slice.append(bbox)
                    flag = False  # flag of true positive or false positive
                    for ilabel in xrange(len(label)):
                        nodule = label[ilabel]
                        if np.linalg.norm((candi_center - nodule['center']) * nodule['spacing']) < nodule['radius']:
                            is_found[ilabel] = 1
                            flag = True
                    if flag:
                        cnt_tp = cnt_tp + 1
                    else:
                        cnt_fp = cnt_fp + 1

                candis[icandi] = _candi_slice

            if not os.path.exists(os.path.join(proj_dir, CANDIS_LUNGMASK_DIR, scan_name)):
                os.mkdir(os.path.join(proj_dir, CANDIS_LUNGMASK_DIR, scan_name))
            np.save(os.path.join(proj_dir, CANDIS_LUNGMASK_DIR, scan_name, 'candis.npy'), candis)
            is_found_all.extend(is_found)
            mali_not_found = [x and (1 - y) for x, y in zip(is_mali, is_found)]
            if np.sum(np.asarray(mali_not_found)) > 0:
                dict_not_found[scan_name] = mali_not_found

            print("%d :%s  >> %4d %4d | %6d %6d %5.2f |  %6d %6d %5.2f | %8d %5.2d | %8d %5.2d"
                  % (icandi,
                      scan_name,
                     np.sum(np.array(is_found) * np.array(is_mali)),
                     np.sum(is_mali),
                     np.sum(np.array(is_found_all)),
                     len(is_found_all),
                     np.sum(np.array(is_found_all)) * 100.0 / len(is_found_all),
                     np.sum(np.array(is_found_all) * np.array(is_mali_all)),
                     np.sum(is_mali_all),
                     np.sum(np.array(is_found_all) * np.array(is_mali_all)) * 100.0 / np.sum(is_mali_all),
                     cnt_tp,
                     cnt_tp * 1.0 / iscan,
                     cnt_fp,
                     cnt_fp * 1.0 / iscan
                     )
                  )

            # np.save(os.path.join(proj_dir, 'output', 'iter_'+CANDIS_DIR.split('_')[1]+'_dict_not_found.npy'), dict_not_found)
    # np.save('is_mali_all.npy', is_mali_all)
    # np.save('is_found_allnpy', is_found_all)
    np.save('not_found.npy', dict_not_found)
    lungmask_kaggle.close()
    lungmask_lidc.close()
    dataset_kaggle.close()
    dataset_lidc.close()



def filter_lungmask_onefold():

    FLAG_LABEL = False
    CANDIS_DIR = '/ssd1/dingj/kaggle/faster_rcnn/all_nodule_flip_ad/output/candis_kaggle_testset_60000'
    CANDIS_OUT_DIR = '/ssd1/dingj/kaggle/faster_rcnn/all_nodule_flip_ad/output/candis_kaggle_testset_before_nms_lungmask_60000'

    if not os.path.exists(CANDIS_OUT_DIR):
        os.mkdir(CANDIS_OUT_DIR)
    scan_names = os.listdir(CANDIS_DIR)
    lungmask_kaggle = h5py.File('/data-disk/Kaggle/backup/kaggle/lungmask.hdf5', 'r')
    dataset_kaggle = h5py.File('/data-disk/Kaggle/backup/kaggle/vol.hdf5', 'r')
    lungmask_spie = h5py.File('/data-disk/Kaggle/backup/spie/lungmask.hdf5', 'r')
    dataset_spie = h5py.File('/data-disk/Kaggle/backup/spie/vol.hdf5', 'r')
    lungmask_lidc = h5py.File('/data-disk/Kaggle/backup/lidc/lungmask.hdf5', 'r')
    dataset_lidc = h5py.File('/data-disk/Kaggle/backup/lidc/vol.hdf5', 'r')

    label_dict = np.load('/data-disk/Kaggle/backup/spie/list/label_dict.npy').item()

    is_found_all = []
    is_mali_all = []
    dict_not_found = {}
    cnt_fp = 0
    cnt_tp = 0

    iscan = 0
    for scan_name in scan_names:
        iscan = iscan+1
        candis = np.load(CANDIS_DIR + '/' + scan_name + '/' + 'candis.npy')
        lungmask = lungmask_kaggle[scan_name][:]
        spacing = dataset_kaggle[scan_name].attrs['spacing']

        if FLAG_LABEL:
            label = label_dict[scan_name]
            is_found = np.zeros((len(label), 1))
            is_mali = np.zeros((len(label), 1))

        for icandi in xrange(len(candis)): # for every slice
            candi_slice = candis[icandi]
            if len(candi_slice) == 0:
                continue

            _candi_slice = []
            for bbox in candi_slice: # for every candidate in one slice
                # if bbox[4] < thresh_val:
                #     continue

                if np.max([(bbox[3]-bbox[1]) , (bbox[2]-bbox[0])] * spacing[0:2]) < 5: # drop bbox with size less than 5mm
                    continue

                candi_r = max([(bbox[3]-bbox[1])/2, (bbox[2]-bbox[0])/2])
                candi_center = [(bbox[1] + bbox[3]) / 2, (bbox[0] + bbox[2]) / 2, icandi]
                if candi_r*spacing[0] < 10:
                    if not lungmask[candi_center[0], candi_center[1], candi_center[2]]:
                        continue

                _candi_slice.append(bbox)
                if FLAG_LABEL:
                    flag = False  # flag of true positive or false positive
                    for ilabel in xrange(len(label)):
                        nodule = label[ilabel]
                        # r = nodule['radius']
                        r = 10
                        if np.linalg.norm((np.asarray(candi_center) - np.asarray(nodule['center'])) * np.asarray(spacing)) < r:
                            is_found[ilabel] = 1
                            flag = True
                    if flag:
                        cnt_tp = cnt_tp + 1
                    else:
                        cnt_fp = cnt_fp + 1

            candis[icandi] = _candi_slice
        if not os.path.exists(CANDIS_OUT_DIR + '/' + scan_name):
            os.makedirs(CANDIS_OUT_DIR + '/' + scan_name)
        np.save(CANDIS_OUT_DIR + '/' + scan_name + '/' + 'candis.npy', candis)

        if FLAG_LABEL:
            is_found_all.extend(is_found)
            mali_not_found = [x and (1 - y) for x, y in zip(is_mali, is_found)]
            if np.sum(np.asarray(mali_not_found)) > 0:
                dict_not_found[scan_name] = mali_not_found
            print("%s  >> %4d %4d | %6d %6d %5.2f |  %6d %6d %5.2f | %8d %5.2d | %8d %5.2d"
                  % (scan_name,
                     np.sum(np.array(is_found) * np.array(is_mali)),
                     np.sum(is_mali),
                     np.sum(np.array(is_found_all)),
                     len(is_found_all),
                     np.sum(np.array(is_found_all)) * 100.0 / len(is_found_all),
                     np.sum(np.array(is_found_all) * np.array(is_mali_all)),
                     np.sum(is_mali_all),
                     np.sum(np.array(is_found_all) * np.array(is_mali_all)) * 100.0 / np.sum(is_mali_all),
                     cnt_tp,
                     cnt_tp * 1.0 / iscan,
                     cnt_fp,
                     cnt_fp * 1.0 / iscan
                     )
                  )
        else:
            print iscan


    lungmask_kaggle.close()
    dataset_kaggle.close()
    lungmask_spie.close()
    dataset_spie.close()
    lungmask_lidc.close()
    dataset_lidc.close()


def cdd_filter_lungmask(thresh_r,
                        candis_dir,
                        candis_lungmask_dir,
                        kaggle_lungmask_path,
                        lidc_lungmask_path,
                        spie_lungmask_path,
                        stage2_lungmask_path,
                        kaggle_vol_path,
                        lidc_vol_path,
                        spie_vol_path,
                        stage2_vol_path):

    if not os.path.exists(candis_dir):
        os.mkdir(candis_dir)
    CANDIS_DIR = candis_dir
    CANDIS_OUT_DIR = candis_lungmask_dir

    if not os.path.exists(CANDIS_OUT_DIR):
        os.makedirs(CANDIS_OUT_DIR)
    scan_names = os.listdir(CANDIS_DIR)
    if kaggle_lungmask_path is not None:
        lungmask_kaggle = h5py.File(kaggle_lungmask_path, 'r')
    else:
        lungmask_kaggle = None
    if kaggle_vol_path is not None:
        dataset_kaggle = h5py.File(kaggle_vol_path, 'r')
    else:
        dataset_kaggle = None
    if spie_lungmask_path is not None:
        lungmask_spie = h5py.File(spie_lungmask_path, 'r')
    else:
        lungmask_spie = None
    if spie_vol_path is not None:
        dataset_spie = h5py.File(spie_vol_path, 'r')
    else:
        dataset_spie = None
    if lidc_lungmask_path is not None:
        lungmask_lidc = h5py.File(lidc_lungmask_path, 'r')
    else:
        lungmask_lidc = None
    if lidc_vol_path is not None:
        dataset_lidc = h5py.File(lidc_vol_path, 'r')
    else:
        dataset_lidc = None
    if stage2_vol_path is not None:
        dataset_stage2 = h5py.File(stage2_vol_path)
    else:
        dataset_stage2 = None
    if stage2_lungmask_path is not None:
        lungmask_stage2 = h5py.File(stage2_lungmask_path)
    else:
        lungmask_stage2 = None

    if dataset_kaggle is not None:
        list_kaggle = dataset_kaggle.keys()
    else:
        list_kaggle = None
    if dataset_lidc is not None:
        list_lidc = dataset_lidc.keys()
    else:
        list_lidc = None
    if dataset_spie is not None:
        list_spie = dataset_spie.keys()
    else:
        list_spie = None
    if stage2_vol_path is not None:
        list_stage2 = dataset_stage2.keys()
    else:
        list_stage2 = None

    iscan = 0
    for scan_name in scan_names:
        iscan = iscan+1
        try:
            candis = np.load(CANDIS_DIR + '/' + scan_name + '/' + 'candis.npy')
        except:
            continue

        if list_lidc is not None and scan_name in list_lidc:
            lungmask = lungmask_lidc[scan_name][:]
        elif list_kaggle is not None and scan_name in list_kaggle:
            lungmask = lungmask_kaggle[scan_name][:]
        elif list_spie is not None and scan_name in list_spie:
            lungmask = lungmask_spie[scan_name][:]
        elif list_stage2 is not None and scan_name in list_stage2:
            lungmask = lungmask_stage2[scan_name][:]
        else:
            print('can not found scan name in dataset')
            exit()
        if list_lidc is not None and scan_name in list_lidc:
            spacing = dataset_lidc[scan_name].attrs['spacing']
        elif list_kaggle is not None and scan_name in list_kaggle:
            spacing = dataset_kaggle[scan_name].attrs['spacing']
        elif list_spie is not None and scan_name in list_spie:
            spacing = dataset_spie[scan_name].attrs['spacing']
        elif list_stage2 is not None and scan_name in list_stage2:
            spacing = dataset_stage2[scan_name].attrs['spacing']
        else:
            print('can not found scan name in dataset')
            exit()

        for icandi in xrange(len(candis)):  # for every slice
            candi_slice = candis[icandi]
            if len(candi_slice) == 0:
                continue
            _candi_slice = []
            for bbox in candi_slice: # for every candidate in one slice
                if np.max([(bbox[3]-bbox[1]) , (bbox[2]-bbox[0])] * spacing[0:2]) < 5:  # drop bbox with size less than 5mm
                    continue
                candi_r = max([(bbox[3]-bbox[1])/2, (bbox[2]-bbox[0])/2])
                candi_center = [(bbox[1] + bbox[3]) / 2, (bbox[0] + bbox[2]) / 2, icandi]
                if candi_r*spacing[0] < thresh_r:
                    if not lungmask[candi_center[0], candi_center[1], candi_center[2]]:
                        continue
                _candi_slice.append(bbox)

            candis[icandi] = _candi_slice
        if not os.path.exists(CANDIS_OUT_DIR + '/' + scan_name):
            os.makedirs(CANDIS_OUT_DIR + '/' + scan_name)
        np.save(CANDIS_OUT_DIR + '/' + scan_name + '/' + 'candis.npy', candis)
        print iscan
    if lungmask_kaggle is not None:
        lungmask_kaggle.close()
    if dataset_kaggle is not None:
        dataset_kaggle.close()
    if lungmask_spie is not None:
        lungmask_spie.close()
    if dataset_spie is not None:
        dataset_spie.close()
    if lungmask_lidc is not None:
        lungmask_lidc.close()
    if dataset_lidc is not None:
        dataset_lidc.close()
    if dataset_stage2 is not None:
        dataset_stage2.close()
    if lungmask_stage2 is not None:
        lungmask_stage2.close()

if __name__ == '__main__':
    # eval_onefold()
    # filter_lungmask()
    filter_lungmask_onefold()